import os
import time
from datetime import timedelta
from typing import Dict, Any, Generator, List, Optional
import easyocr
from openai import OpenAI, APIError
from PyPDF2 import PdfReader
from docx import Document
import numpy as np
import cv2


def classify_injury_assessment_type(text: str) -> str:
    """根据文本内容识别工伤鉴定书类型"""
    categories = {
        # 按伤残等级
        "一级伤残鉴定书": ["一级伤残", "伤残等级：一级", "伤残程度：一级"],
        "二级伤残鉴定书": ["二级伤残", "伤残等级：二级", "伤残程度：二级"],
        "三级伤残鉴定书": ["三级伤残", "伤残等级：三级", "伤残程度：三级"],
        "四级伤残鉴定书": ["四级伤残", "伤残等级：四级", "伤残程度：四级"],
        "五级伤残鉴定书": ["五级伤残", "伤残等级：五级", "伤残程度：五级"],
        "六级伤残鉴定书": ["六级伤残", "伤残等级：六级", "伤残程度：六级"],
        "七级伤残鉴定书": ["七级伤残", "伤残等级：七级", "伤残程度：七级"],
        "八级伤残鉴定书": ["八级伤残", "伤残等级：八级", "伤残程度：八级"],
        "九级伤残鉴定书": ["九级伤残", "伤残等级：九级", "伤残程度：九级"],
        "十级伤残鉴定书": ["十级伤残", "伤残等级：十级", "伤残程度：十级"],
        
        # 按鉴定类型
        "劳动能力鉴定书": ["劳动能力鉴定", "劳动功能障碍程度", "生活自理障碍程度"],
        "工伤认定书": ["工伤认定", "认定为工伤", "工伤认定决定书"],
        "职业病诊断书": ["职业病诊断", "职业病", "职业性疾病"],
        "医疗终结鉴定书": ["医疗终结", "治疗终结", "医疗期满"],
        "复查鉴定书": ["复查鉴定", "重新鉴定", "再次鉴定"],
        
        # 按伤害部位
        "头部外伤鉴定书": ["头部外伤", "颅脑损伤", "脑外伤"],
        "脊柱损伤鉴定书": ["脊柱损伤", "脊椎骨折", "腰椎", "颈椎"],
        "四肢损伤鉴定书": ["四肢损伤", "手部损伤", "足部损伤", "上肢", "下肢"],
        "内脏损伤鉴定书": ["内脏损伤", "胸腹部损伤", "内脏破裂"],
        "烧伤鉴定书": ["烧伤", "烫伤", "化学烧伤"],
        "听力损伤鉴定书": ["听力损失", "耳聋", "听力障碍"],
        "视力损伤鉴定书": ["视力损失", "失明", "视力障碍"],
        
        # 按鉴定机构
        "省级鉴定书": ["省劳动能力鉴定委员会", "省级鉴定"],
        "市级鉴定书": ["市劳动能力鉴定委员会", "地市级鉴定"],
        "司法鉴定书": ["司法鉴定", "司法鉴定所", "司法鉴定中心"]
    }

    for name, keywords in categories.items():
        if any(keyword in text for keyword in keywords):
            return name
    return "无法识别/可能为非标准工伤鉴定书"


def detect_official_seal_in_text(text: str) -> bool:
    """通过关键字判断OCR文本中是否有官方印章相关字样"""
    seal_keywords = ["劳动能力鉴定委员会", "鉴定专用章", "公章", "盖章", "（章）", 
                    "鉴定机构章", "医疗机构章", "司法鉴定章"]
    return any(k in text for k in seal_keywords)


def detect_seal_in_image(file_path: str, min_area_ratio: float = 0.001) -> bool:
    """检测图像中是否有红色圆形印章"""
    if not os.path.exists(file_path):
        return False
    arr = np.fromfile(file_path, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        return False
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    # 红色在HSV空间范围
    lower1, upper1 = np.array([0, 70, 50]), np.array([10, 255, 255])
    lower2, upper2 = np.array([170, 70, 50]), np.array([180, 255, 255])
    mask = cv2.inRange(hsv, lower1, upper1) | cv2.inRange(hsv, lower2, upper2)
    # 计算红色像素比例
    red_ratio = np.sum(mask > 0) / (mask.size + 1e-6)
    if red_ratio < min_area_ratio:
        return False
    # 模糊后找圆
    blurred = cv2.GaussianBlur(mask, (9, 9), 2)
    circles = cv2.HoughCircles(
        blurred, cv2.HOUGH_GRADIENT, dp=1.2, minDist=100,
        param1=50, param2=35, minRadius=30, maxRadius=150
    )
    return circles is not None


class InjuryAssessmentAnalyzer:
    def __init__(self, api_key: str):
        self.client = OpenAI(
            api_key=api_key,
            base_url="https://api.deepseek.com"
        )

    def _extract_from_image(self, file_path: str, reader=None) -> str:
        try:
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在: {file_path}")
            img_array = np.fromfile(file_path, dtype=np.uint8)
            image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
            if image is None:
                raise ValueError("无法解析图像，cv2.imdecode 失败")
            reader = reader or easyocr.Reader(['ch_sim', 'en'])
            result = reader.readtext(image, detail=0)
            return "\n".join(result).strip()
        except Exception as e:
            raise RuntimeError(f"图片OCR识别失败: {str(e)}")

    def extract_text(self, file_path: str) -> str:
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        ext = file_path.lower()
        if ext.endswith('.pdf'):
            return self._extract_from_pdf(file_path)
        elif ext.endswith('.docx'):
            return self._extract_from_docx(file_path)
        elif ext.endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.webp')):
            return self._extract_from_image(file_path)
        else:
            raise ValueError("不支持的文件格式，请提供PDF、DOCX或图片文件")

    def extract_texts_from_multiple_images(self, image_paths: List[str]) -> str:
        reader = easyocr.Reader(['ch_sim', 'en'])
        all_text = []
        for path in image_paths:
            try:
                text = self._extract_from_image(path, reader)
                all_text.append(text)
            except Exception as e:
                print(f"\n⚠️ 图片处理失败 [{path}]: {str(e)}")
        return "\n\n".join(all_text)

    def _extract_from_pdf(self, file_path: str) -> str:
        text = ""
        try:
            with open(file_path, 'rb') as file:
                reader = PdfReader(file)
                for page in reader.pages:
                    text += page.extract_text() or ""
        except Exception as e:
            raise RuntimeError(f"PDF文件读取失败: {str(e)}")
        return text

    def _extract_from_docx(self, file_path: str) -> str:
        try:
            doc = Document(file_path)
            return "\n".join([para.text for para in doc.paragraphs if para.text])
        except Exception as e:
            raise RuntimeError(f"DOCX文件读取失败: {str(e)}")

    def analyze_injury_assessment_stream(
        self,
        text: str,
        assessment_type: str,
        file_paths: Optional[List[str]] = None
    ) -> Generator[str, None, Dict[str, Any]]:
        if not text.strip():
            raise ValueError("工伤鉴定书文本内容为空")

        # 文本层面印章检测
        has_seal_text = detect_official_seal_in_text(text)
        # 图像层面印章检测
        has_seal_image = False
        if file_paths:
            for p in file_paths:
                if p.lower().endswith(('.png','.jpg','.jpeg','.bmp','.tiff')) and detect_seal_in_image(p):
                    has_seal_image = True
                    break

        seal_note = "（检测到官方印章）" if has_seal_text and has_seal_image else "（未检测到官方印章）"
        print(seal_note)

        system_prompt = f"""你是一位专业的工伤法律师，负责分析工伤鉴定书的合规性和法律效力。
该鉴定书初步识别为：{assessment_type} {seal_note}

请严格按照以下要求进行分析：

1. **鉴定书类型确认**：
   - 确认鉴定书的具体类型（工伤认定书/劳动能力鉴定书/职业病诊断书等）
   - 判断伤残等级（如适用）
   - 评估鉴定书的规范性和完整性

2. **鉴定机构和程序合规性分析**：
   - 鉴定机构是否具备相应资质
   - 鉴定程序是否符合法定要求
   - 鉴定时间是否在法定期限内
   - 是否有相关印章和签字

3. **鉴定内容专业性评估**：
   - 伤残等级认定是否合理
   - 劳动功能障碍程度评定
   - 生活自理障碍程度评定
   - 医疗依据是否充分

4. **法律效力分析**：
   - 该鉴定书的法律效力等级
   - 是否可以作为赔偿依据
   - 对后续维权的影响

5. **赔偿和待遇分析**：
   - 根据伤残等级可享受的工伤保险待遇
   - 一次性伤残补助金标准
   - 伤残津贴标准（如适用）
   - 医疗费用报销范围
   - 护理费、康复费等相关费用

6. **维权建议**：
   - 如对鉴定结果不服的救济途径
   - 申请复查或重新鉴定的条件和程序
   - 后续法律维权的策略建议
   - 需要收集的补充证据

7. **风险提示**：
   - 鉴定书可能存在的问题或瑕疵
   - 时效性要求和注意事项
   - 可能影响赔偿的因素

请提供专业、详细的法律分析，重点关注当事人的合法权益保护。"""

        try:
            stream = self.client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"请分析以下工伤鉴定书：\n{text[:15000]}"}
                ],
                temperature=0.3,
                max_tokens=3000,  # 工伤鉴定分析需要更详细的输出
                stream=True
            )

            collected_content = []
            total_tokens = 0
            for chunk in stream:
                if not chunk.choices:
                    continue
                delta = chunk.choices[0].delta
                if delta and delta.content:
                    collected_content.append(delta.content)
                    yield delta.content
                if hasattr(chunk, 'usage') and chunk.usage:
                    total_tokens = chunk.usage.total_tokens

            return {"metadata": {"total_tokens": total_tokens, "complete_response": "".join(collected_content)}}

        except APIError as e:
            raise RuntimeError(f"API请求失败: {str(e)}")
        except Exception as e:
            raise RuntimeError(f"分析过程中出错: {str(e)}")


def main():
    API_KEY = "sk-20856422ed6644e3827b9d5403c9542a"  # 替换为你的API密钥
    analyzer = InjuryAssessmentAnalyzer(API_KEY)

    print("工伤鉴定书分析工具（流式输出版）")
    print("=" * 40)
    file_input = input("请输入工伤鉴定书路径（多个图片用英文逗号分隔，或一个PDF/DOCX）: ").strip()
    file_paths = [p.strip() for p in file_input.split(',') if p.strip()]

    if not file_paths:
        print("❌ 未输入有效路径")
        return

    # 判断模式
    if len(file_paths) > 1:
        if all(p.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff')) for p in file_paths):
            mode = "multi_image"
        else:
            print("❌ 当前仅支持：多张图片 或 单个 PDF/DOCX。请检查输入。")
            return
    elif file_paths[0].lower().endswith(('.pdf', '.docx')):
        mode = "single_document"
    else:
        print("❌ 当前仅支持：多张图片 或 单个 PDF/DOCX。请检查输入。")
        return

    try:
        # 文本提取
        if mode == "multi_image":
            text = analyzer.extract_texts_from_multiple_images(file_paths)
        else:
            text = analyzer.extract_text(file_paths[0])

        if not text.strip():
            print("❌ 没有提取到有效文本，终止分析")
            return

        # 初步识别类型
        assessment_type = classify_injury_assessment_type(text)
        print(f"\n📌 初步识别的鉴定书类型：{assessment_type}")

        print("\n正在分析工伤鉴定书，请稍候...\n")
        print("=" * 40)
        print("实时分析结果:")
        print("=" * 40)

        start_time = time.time()
        full_response = []
        metadata = {}

        # 分析并输出
        for chunk in analyzer.analyze_injury_assessment_stream(text, assessment_type, file_paths):
            print(chunk, end="", flush=True)
            full_response.append(chunk)
        elapsed = time.time() - start_time

        print("\n\n" + "=" * 40)
        print("分析完成!")
        if metadata and metadata.get('total_tokens', 0) > 0:
            print(f"\n总Tokens: {metadata['total_tokens']}")
            print(f"总耗时: {timedelta(seconds=elapsed)}")
            print(f"处理速度: {metadata['total_tokens']/elapsed:.2f} tokens/秒")

        # 保存结果
        save_path = os.path.join(os.path.dirname(file_paths[0]), "injury_assessment_analysis_result.txt")
        with open(save_path, 'w', encoding='utf-8') as f:
            f.write("".join(full_response))
        print(f"\n✅ 分析结果已保存到: {save_path}")

    except KeyboardInterrupt:
        print("\n用户中断操作")
    except Exception as e:
        print(f"\n程序发生错误: {str(e)}")


if __name__ == "__main__":
    main()